#include "PathPlan.h"
#include <stdlib.h>

using namespace common::planning;

/**
 * @brief 对start到goal的直线段路径, 按间距resolution进行采样/插值
 */
PoseDq common::planning::InterpolatePoses(const Pose& start, 
                                          const Pose& goal,
                                          const double resolution) {
  PoseDq interpolated_poses{start};
  auto len = start.EuclidDistance(goal);
  double cnt_d = len / resolution;
  int cnt = static_cast<int>(cnt_d);
  auto dx = cnt > 0 ? (goal.X() - start.X()) / cnt_d : 0.0;
  auto dy = cnt > 0 ? (goal.Y() - start.Y()) / cnt_d : 0.0;
  Pose pose = start;
  for (int k = 1; k <= cnt; ++k) {
    pose.SetX(pose.X() + dx);
    pose.SetY(pose.Y() + dy);
    interpolated_poses.emplace_back(pose);
  }
  if (interpolated_poses.size() > 1 && // 返回结果要包括start和goal
      interpolated_poses.back().EuclidDistance(goal) < 0.5 * resolution) {
    interpolated_poses.back() = goal;
  } else {
    interpolated_poses.emplace_back(goal);
  }
  return interpolated_poses;
}

/**
 * @brief 把result_info(PathPlan::ResultType)转换成字符串
 */
std::string common::planning::ToStringResult(const uint8_t result_info) {
  std::string result;
  switch (result_info) {
    case PathPlan::START_TRAPPED:
      result = "Path plan fail, start trapped !";
      break;
    case PathPlan::GOAL_TRAPPED:
      result = "Path plan fail, goal trapped !";
      break;
    case PathPlan::START_GOAL_NOT_CONNECTED:
      result = "Path plan fail, start goal not connected !";
      break;
    case PathPlan::SEARCH_FAIL:
      result = "Path plan fail, for unexpected reason !";
      break;
    default:
      result = "Path plan succeed !";
      break;
  }
  return result;
}

/**
 * @brief 初始化
 * @param binary_map 二值化的栅格地图（0-障碍物, 非0-空闲)
 * @param origin 栅格图左下角栅格中心对应的物理坐标
 * @param width 栅格图的宽, unit: 栅格
 * @param height 栅格图的高, unit: 栅格
 * @param resolution 栅格图的分辨率, unit: m
 * @param bump_dis 碰撞距离, unit: m (最近障碍物距离小于等于该值的栅格, 视为不可行区域)
 * @param safe_dis 安全距离, unit: m (最近障碍物距离大于该值的栅格, 视为足够安全的区域)
 */
void PathPlan::Init(const std::vector<uint8_t>& binary_map,
                    const Pose& origin, 
                    const int64_t width,
                    const int64_t height, 
                    const double resolution,
                    const double bump_dis,
                    const double safe_dis) {
  // 初始化障碍物距离场信息DFMap()
  DFMap().SetOrigin(origin);
  DFMap().SetWidth(width);
  DFMap().SetHeight(height);
  DFMap().SetResolution(resolution);
  DFMap().MapData().resize(binary_map.size(), 0);
  auto bump_grid_dis = bump_dis / resolution;
  auto safe_grid_dis = safe_dis / resolution;
  DFMap().SetObstacleFreeThresh(bump_grid_dis * bump_grid_dis, safe_grid_dis * safe_grid_dis);
  LOG(INFO) << "bump_dis is " << bump_dis << ", safe_dis is " << safe_dis << ", binary_map.size " << binary_map.size();
  
  // 构造障碍物距离场df_mat、DFMap().MapData(), 根据bump_dis二值化得到configuration_space
  cv::Mat binary_mat = cv::Mat::zeros(height, width, CV_8UC1);
  size_t idx = 0;
  for (int64_t iy = 0; iy < height && idx < binary_map.size(); ++iy) {
    for (int64_t ix = 0; ix < width && idx < binary_map.size(); ++ix) {
      binary_mat.ptr<uchar>(iy)[ix] = binary_map.at(idx++);
    }
  }
  cv::Mat configuration_space = cv::Mat::ones(height, width, CV_8UC1); // 配置空间 (二值化, 非0用于提取连通域)
  cv::Mat df_mat;
  cv::distanceTransform(binary_mat, df_mat, cv::DIST_L2, cv::DIST_MASK_3);
  for (int64_t iy = 0; iy < height; ++iy) {
    for (int64_t ix = 0; ix < width; ++ix) {
      auto& obstacle_dist = df_mat.ptr<float>(iy)[ix];
      DFMap().SetGrid(ix, iy, pow(obstacle_dist, 2)); // 平方距离
      if (obstacle_dist <= bump_grid_dis) {
        configuration_space.ptr<uchar>(iy)[ix] = 0;
      }
    }
  }

  // 标记连通域connected_
  connected_ = cv::Mat::zeros(configuration_space.size(), CV_32S); // 元素: 32-bit signed
  int num_connected_ = cv::connectedComponents(configuration_space, connected_, 8, CV_32S); // 连通域个数
  LOG(INFO) << "num_connected_ is " << num_connected_;

  // 调试用
  // bool test_gpp = ros::param::param("test_gpp", false);
  // if (test_gpp) {
    std::vector<cv::Vec3b> colors(num_connected_ + 1);
    colors[0] = cv::Vec3b(0, 0, 0);  //背景色
    for (int m = 1; m <= num_connected_; m++) {
      colors[m] = cv::Vec3b(rand() % 256, rand() % 256, rand() % 256); // 随机色
    }
    cv::Mat tmp_map(height, width, CV_8UC3);
    for (int64_t iy = 0; iy < height; ++iy) {
      for (int64_t ix = 0; ix < width; ++ix) {
        auto tmp_label = connected_.ptr<int32_t>(iy)[ix];
        tmp_map.at<cv::Vec3b>(iy, ix) = colors[tmp_label];
      }
    }
    cv::flip(tmp_map, tmp_map, 0);
    auto map_file = ::common::log::kGlobalPathPlannerLogPath + "/" + common::log::CurTimeString() + "_connected_regions.png";
    cv::imwrite(map_file, tmp_map);
  // }
}

/**
 * @brief 搜索最小通行代价路径
 * @param start 起点, 物理坐标
 * @param goal 目标点, 物理坐标
 * @param resolution 搜索得到的路径点间距, unit: m
 * @param method 图搜算法(GraphSearch::DIJKSTRA, GraphSearch::ASTAR, GraphSearch::JPS)
 * @return PoseDq 路径, 物理坐标
 */
PoseDq PathPlan::SearchPath(const Pose& start, 
                            const Pose& goal, 
                            const double resolution,
                            const uint8_t method) {
  Node start_node(DFMap().X2IX(start.X()), DFMap().Y2IY(start.Y()));
  Node goal_node(DFMap().X2IX(goal.X()), DFMap().Y2IY(goal.Y()));
  bool plan_result = g_search_.Plan(start_node, goal_node, DFMap().ThreshObstacle(), DFMap().ThreshFree(), method);
  
  // 栅格坐标路径 - 关键点
  Nodes search_path{};
  if (plan_result) {
    search_path = g_search_.GetOptimizedPath();
  } else {
    LOG(WARNING) << "g_search_ failed ...";
  }

  // 物理坐标路径: 关键路点 -> 插值
  PoseDq path{};
  if(search_path.size() >= 2) { // 规划成功 -> 该条件为true
    path.emplace_back(start);
    // 把损失精度的头尾去掉
    search_path.pop_front();
    search_path.pop_back();
    for (const auto& node : search_path) {
      path.emplace_back(Pose(DFMap().IX2X(node.x), DFMap().IY2Y(node.y)));
    }
    path.emplace_back(goal);

    // 按给定间距插值
    PoseDq sampled_poses{};
    for (size_t i = 0; i+1 < path.size(); ++ i) {
      auto poses = InterpolatePoses(path.at(i), path.at(i+1), resolution);
      if (i+1 != path.size()) {
        poses.pop_back(); // 防止插入重复点
      }
      sampled_poses.insert(sampled_poses.end(), poses.begin(), poses.end());
    }

    path.swap(sampled_poses);
  }

  return path; // 关键点路径, 规划失败则为空
}

/**
 * @brief pose对应的栅格是否被占用
 */
bool PathPlan::IsOccupied(const Pose& pose) {
  auto ix = DFMap().X2IX(pose.X());
  auto iy = DFMap().Y2IY(pose.Y());
  return (DFMap().Occupied(ix, iy));
}

/**
 * @brief 在pose周围range(unit: m)范围内找到一个可行点, 更新pose为该可行点
 * @return true 更新成功
 * @return false range(unit: m)范围内无可行点
 */
bool PathPlan::UpdateToFree(Pose& pose, const double range) {
  int64_t range_i = range / DFMap().Resolution(); // 栅格范围
  bool updated = false;

  // 8-邻域bfs搜索
  auto seed_ix = DFMap().X2IX(pose.X());
  auto seed_iy = DFMap().Y2IY(pose.Y());
  CellDq visited_q{Cell(seed_ix, seed_iy)};
  CellSet expanded_s{Cell(seed_ix, seed_iy)};
  while(!visited_q.empty()) {
    auto pop_grid = visited_q.front();
    visited_q.pop_front();
    if (!DFMap().Occupied(pop_grid.X(), pop_grid.Y())) {
      if (seed_ix != pop_grid.X() || seed_iy != pop_grid.Y()) { // 种子点不需要更新
        pose.SetX(DFMap().IX2X(pop_grid.X()));
        pose.SetY(DFMap().IY2Y(pop_grid.Y()));
      }
      updated = true;
      break;
    }
    for (int64_t i = -1; i <= 1; ++i) {
      for (int64_t j = -1; j <= 1; ++j) {
        if (i == 0 && j == 0) {
          continue;
        }
        // 8-邻域
        Cell grid(pop_grid.X() + i, pop_grid.Y() + j);
        if (expanded_s.find(grid) == expanded_s.end() &&
            std::hypot(grid.X() - seed_ix, grid.Y() - seed_iy) <= range_i) {  // range_i范围内
          visited_q.push_back(grid);
          expanded_s.insert(grid);
        }
      }
    }
  }
  return updated;
}

/**
 * @brief 获取栅格(ix, iy)的连通图标记值
 */
int32_t PathPlan::GetConnectedValue(const int64_t ix, const int64_t iy) {
  return DFMap().Occupied(ix, iy) ? 0 : connected_.ptr<int32_t>(iy)[ix];
}

/**
 * @brief 在pose周围range(unit: m)范围内找到一个连通图上标记为connected_value的可行点, 更新pose为该可行点
 * @return true 更新成功
 * @return false range(unit: m)范围内无符合条件的可行点
 */
bool PathPlan::UpdateToConnected(Pose& pose, const double range, const int32_t connected_value) {
  int64_t range_i = range / DFMap().Resolution(); // 栅格范围
  bool updated = false;

  // 8-邻域bfs搜索
  auto seed_ix = DFMap().X2IX(pose.X());
  auto seed_iy = DFMap().Y2IY(pose.Y());
  CellDq visited_q{Cell(seed_ix, seed_iy)};
  CellSet expanded_s{Cell(seed_ix, seed_iy)};
  while(!visited_q.empty()) {
    auto pop_grid = visited_q.front();
    visited_q.pop_front();
    if (GetConnectedValue(pop_grid.X(), pop_grid.Y()) == connected_value) {
      if (seed_ix != pop_grid.X() || seed_iy != pop_grid.Y()) { // 种子点不需要更新
        pose.SetX(DFMap().IX2X(pop_grid.X()));
        pose.SetY(DFMap().IY2Y(pop_grid.Y()));
      }
      updated = true;
      break;
    }
    for (int64_t i = -1; i <= 1; ++i) {
      for (int64_t j = -1; j <= 1; ++j) {
        if (i == 0 && j == 0) {
          continue;
        }
        // 8-邻域
        Cell grid(pop_grid.X() + i, pop_grid.Y() + j);
        if (expanded_s.find(grid) == expanded_s.end() &&
            std::hypot(grid.X() - seed_ix, grid.Y() - seed_iy) <= range_i) {  // range_i范围内
          visited_q.push_back(grid);
          expanded_s.insert(grid);
        }
      }
    }
  }
  return updated;
}

/**
 * @brief 检查p1和p2是否连通
 */
bool PathPlan::IsConnected(const Pose& p1, const Pose& p2) {
  auto p1_ix = DFMap().X2IX(p1.X());
  auto p1_iy = DFMap().Y2IY(p1.Y());
  auto p1_value = GetConnectedValue(p1_ix, p1_iy);
  if (p1_value == 0) {
    return false;
  }
  auto p2_ix = DFMap().X2IX(p2.X());
  auto p2_iy = DFMap().Y2IY(p2.Y());
  auto p2_value = GetConnectedValue(p2_ix, p2_iy);
  if (p2_value == 0) {
    return false;
  }

  return p1_value == p2_value;
}

/**
 * @brief 搜索最小通行代价路径
 * @param start 起点, 物理坐标
 * @param start_update_range 起点如果被占用, 允许在多大范围内找可行点更新起点, unit: m
 * @param goal 目标点, 物理坐标
 * @param goal_update_range 目标点如果被占用, 允许在多大范围内找可行点更新目标点, unit: m
 * @param resolution 搜索得到的路径点间距, unit: m
 * @param result_info 调用方可据此获取规划结果的详细信息 (取值范围: PathPlan::ResultType)
 * @param method 图搜算法(GraphSearch::DIJKSTRA, GraphSearch::ASTAR, GraphSearch::JPS)
 * @return PoseDq 路径, 物理坐标
 */
PoseDq PathPlan::SearchPathWithResult(const Pose& start, 
                                      const double start_update_range, 
                                      const Pose& goal,
                                      const double goal_update_range, 
                                      const double resolution,
                                      uint8_t& result_info, 
                                      const uint8_t method) {
  // 1. 检查、更新起点
  bool start_occupied = IsOccupied(start);
  Pose updated_start = start;
  bool start_updated = UpdateToFree(updated_start, start_update_range);
  if (start_occupied && !start_updated) { // 起点被困, 返回搜索失败
    result_info = PathPlan::START_TRAPPED;
    return PoseDq{};
  }

  // 2. 检查、更新目标点
  bool goal_occupied = IsOccupied(goal);
  Pose updated_goal = goal;
  bool goal_updated = UpdateToFree(updated_goal, goal_update_range);
  if (goal_occupied && !goal_updated) { // 目标点被困, 返回搜索失败
    result_info = PathPlan::GOAL_TRAPPED;
    return PoseDq{};
  }

  // 3. 检查updated_start、updated_goal的连通性
  if (!IsConnected(updated_start, updated_goal)) { // 起点、目标点不连通, 返回搜索失败
    result_info = PathPlan::GOAL_TRAPPED;
    return PoseDq{};
  }

  // 4. 图搜索
  PoseDq tmp_path = SearchPath(updated_start, updated_goal, resolution, method);
  if (tmp_path.empty()) { // 图搜索失败, 返回搜索失败
    result_info = PathPlan::START_GOAL_NOT_CONNECTED;
    return PoseDq{};
  }

  // 5. 图搜索成功, 更新路径
  PoseDq result_path{};

  // 起点更新了, 补充start到updated_start之间的直线路径点
  if (start_occupied && start_updated) {
    LOG(INFO) << "start updated, interpolate from start to updated_start";
    auto poses = InterpolatePoses(start, tmp_path.front(), resolution);
    poses.pop_back(); // 防止插入重复点
    result_path.insert(result_path.end(), poses.begin(), poses.end());
  }

  // updated_start到updated_goal之间的最小通行代价路径
  result_path.insert(result_path.end(), tmp_path.begin(), tmp_path.end());

  // 目标点更新了, 补充updated_goal到goal之间的直线路径点
  if (goal_occupied && goal_updated) {
    LOG(INFO) << "goal updated, interpolate from updated_goal to goal";
    auto poses = InterpolatePoses(tmp_path.back(), goal, resolution);
    poses.pop_front(); // 防止插入重复点
    result_path.insert(result_path.end(), poses.begin(), poses.end());
  }

  result_info = PathPlan::SUCCEED;

  return result_path;
}